#include <iostream>
#include <string>
#include <cstdint>
#include "RSA.h"

using namespace std;

void printHelp()
{
	cout << "0: set rsa parameter" << endl
		<< "1: set public key" << endl
		<< "2: generate public key" << endl
		<< "3: encrypt" << endl
		<< "4: decrypt" << endl
		<< "5: show rsa state" << endl
		<< "q: exit" << endl;
}

template<class T>
void setRSAParam(RSAKeyPair<T>& rsa)
{
	cout << "input two primes: ";
	T p, q;
	cin >> p >> q;
	int result = rsa.setParam(p, q);
	if (result != 0)
		cout << "Fail. The number " << -result << " is not prime." << endl;
	else
		cout << "Success. mod is " << rsa.getN() << endl;
}

template<class T>
void setRSAKey(RSAKeyPair<T>& rsa)
{
	cout << "input public key: ";
	T e;
	cin >> e;
	int result = rsa.setKey(e);
	if (result == rsa.NOPARAM)
		cout << "Fail. Param unset." << endl;
	else if (result != 0)
		cout << "Fail. Illegal public key!" << endl;
	else
		cout << "Success. private key is: " << rsa.getD() << endl;
}

template<class T>
void genRSAKey(RSAKeyPair<T>& rsa)
{
	int result = rsa.setKey();
	if (result == rsa.NOPARAM)
		cout << "Fail. Param unset." << endl;
	else if (result != 0)
		cout << "Fail. Unknown error." << endl;
	else
		cout << "Success. public key is: " << rsa.getE() << endl
		<< "private key is: " << rsa.getD() << endl;
}

template<class T>
void RSAEncrypt(RSAKeyPair<T>& rsa)
{
	if (rsa.getState() != rsa.ENABLE)
	{
		cout << "RSA disabled." << endl;
		return;
	}
	cout << "input message: ";
	T m, c;
	cin >> m;
	int result = rsa.encrypt(m, c);
	if (result != 0)
		cout << "Fail. Unknown error." << endl;
	else
		cout << "cipher: " << c << endl;
}

template<class T>
void RSADecrypt(RSAKeyPair<T>& rsa)
{
	if (rsa.getState() != rsa.ENABLE)
	{
		cout << "RSA disabled." << endl;
		return;
	}
	cout << "input cipher: ";
	T c, m;
	cin >> c;
	int result = rsa.decrypt(c, m);
	if (result != 0)
		cout << "Fail. Unknown error." << endl;
	else
		cout << "message: " << m << endl;
}
template<class T>
void displayRSAState(RSAKeyPair<T>& rsa)
{
	switch (rsa.getState())
	{
	case rsa.ENABLE:
		cout << "All ready.\n n: " << rsa.getN() << endl
			<< " e: " << rsa.getE() << endl
			<< " d: " << rsa.getD() << endl;
		break;
	case rsa.NOKEY:
		cout << "No key\n n: " << rsa.getN() << endl;
		break;
	case rsa.NOPARAM:
		cout << "Uninit" << endl;
	default:
		break;
	}
}

int RSAOperation()
{
	cout << "input '?' to show menu" << endl;
	string input;
	RSAKeyPair<int64_t> rsa;
	bool flag = true;
	while (flag)
	{
		cout << "> ";
		cin >> input;
		switch (input[0])
		{
		case '?':
			printHelp();
			break;
		case '0':
			setRSAParam(rsa);
			break;
		case '1':
			setRSAKey(rsa);
			break;
		case '2':
			genRSAKey(rsa);
			break;
		case '3':
			RSAEncrypt(rsa);
			break;
		case '4':
			RSADecrypt(rsa);
			break;
		case '5':
			displayRSAState(rsa);
			break;
		case 'q':
			flag = false;
			break;
		default:
			cout << "undefined input" << endl;
			break;
		}
	}
	return 0;
}

int main()
{
	RSAOperation();
}